import java.util.*;

public class TopK {


    // 找前 k个最小的元素
    public static int[] topK(int[] arr,int k){
        // 创建一个大小为 k的大根堆
        PriorityQueue<Integer> maxHeap = new PriorityQueue<>(k,new Comparator<Integer>() {
            @Override
            public int compare(Integer o1, Integer o2) {
                return o2 - o1;
            }
        });
        for (int i = 0; i < arr.length; i++) {
            if (i < k){
                // 放入前 k 个元素
                maxHeap.offer(arr[i]);
            }else{
                // 从第 k+1个元素开始进行判断是否要入堆
                if (maxHeap.peek() > arr[i]){
                    maxHeap.poll();
                    maxHeap.offer(arr[i]);
                }
            }
        }
        int[] ret = new int[k];
        for (int i = 0; i < k; i++) {
            ret[i] = maxHeap.poll();
        }
        return ret;
    }


    // 查找和最小的k对数字  1
//    public static List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {
//        PriorityQueue<List<Integer>> maxHeap = new PriorityQueue<>(k, new Comparator<List<Integer>>() {
//            @Override
//            public int compare(List<Integer> o1, List<Integer> o2) {
//                return (o2.get(0)+o2.get(1)) - (o1.get(0)+o1.get(1));
//            }
//        });
////        for (int i = 0; i < Math.min(nums1.length,k); i++) {
////            for (int j = 0; j < Math.min(nums2.length,k); j++) {
//        for (int i = 0; i < nums1.length; i++) {
//            for (int j = 0; j < nums2.length; j++) {
//                if (maxHeap.size() < k){
//                    List<Integer> list = new ArrayList<>();
//                    list.add(nums1[i]);
//                    list.add(nums2[j]);
//                    maxHeap.offer(list);
//                }else{
//                    int top = maxHeap.peek().get(0) + maxHeap.peek().get(1);
//                    if (top > nums1[i] + nums2[j]){
//                        maxHeap.poll();
//                        List<Integer> list = new ArrayList<>();
//                        list.add(nums1[i]);
//                        list.add(nums2[j]);
//                        maxHeap.offer(list);
//                    }
//                }
//            }
//        }
//        List<List<Integer>> ret = new ArrayList<>();
//        for (int i = 0; i < k; i++) {
//            if (!maxHeap.isEmpty()){
//                ret.add(0,maxHeap.poll());
//            }
//        }
//        return ret;
//    }

    // 查找和最小的 k对数字  2
    public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {
        PriorityQueue<int[]> pq = new PriorityQueue<>(new Comparator<int[]>() {
            @Override
            public int compare(int[] o1, int[] o2) {
                return nums1[o1[0]]+nums2[o1[1]]-nums1[o2[0]]-nums2[o2[1]];
            }
        });
        for (int i = 0; i < Math.min(nums1.length,k); i++) {
            pq.offer(new int[]{i,0});
        }
        List<List<Integer>> ret = new ArrayList<>();
        while (k-- > 0 && !pq.isEmpty()){
            int[] top = pq.poll();
            ret.add(Arrays.asList(nums1[top[0]],nums2[top[1]]));
            if (++top[1] < nums2.length){
                pq.offer(top);
            }
        }
        return ret;
    }

    public static void main(String[] args) {
//        int[] arr = {1,3,576,3,85,4};
//        int[] nums = topK(arr, 5);

        // 默认得到一个小根堆
        PriorityQueue<Integer> smallHeap = new PriorityQueue<>();
        smallHeap.offer(23);
        smallHeap.offer(2);
        smallHeap.offer(11);
        System.out.println(smallHeap.poll());// 弹出2，剩余最小的元素就是11，会被调整到堆顶，下一次弹出
        System.out.println(smallHeap.poll());// 弹出11

        // 如果需要得到大根堆，在里面传一个比较器
        PriorityQueue<Integer> BigHeap = new PriorityQueue<>(new Comparator<Integer>() {
            @Override
            public int compare(Integer o1, Integer o2) {
                return o2 - o1;
            }
        });
    }
}
